//! 追踪器实现
//! 
//! 提供分布式追踪器的具体实现，包括Mock实现和生产级实现

use crate::*;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::RwLock;
use async_trait::async_trait;
use uuid::Uuid;

/// Mock分布式追踪器
/// 
/// 用于测试和开发的Mock实现，提供完整的追踪功能模拟
pub struct MockDistributedTracer {
    /// 活跃的spans
    active_spans: Arc<RwLock<HashMap<String, TraceSpan>>>,
    /// 已完成的spans
    finished_spans: Arc<RwLock<Vec<TraceSpan>>>,
    /// 统计信息
    statistics: Arc<RwLock<TraceStatistics>>,
    /// 是否启用
    enabled: bool,
    /// 生成trace_id的计数器
    trace_counter: Arc<RwLock<u32>>,
    /// 生成span_id的计数器
    span_counter: Arc<RwLock<u32>>,
}

impl MockDistributedTracer {
    /// 创建新的Mock追踪器
    pub fn new() -> Self {
        Self {
            active_spans: Arc::new(RwLock::new(HashMap::new())),
            finished_spans: Arc::new(RwLock::new(Vec::new())),
            statistics: Arc::new(RwLock::new(TraceStatistics {
                total_spans: 0,
                active_spans: 0,
                finished_spans: 0,
                avg_span_duration_ms: 0.0,
                sampling_rate: 1.0,
            })),
            enabled: true,
            trace_counter: Arc::new(RwLock::new(0)),
            span_counter: Arc::new(RwLock::new(0)),
        }
    }

    /// 创建禁用的Mock追踪器
    pub fn new_disabled() -> Self {
        let mut tracer = Self::new();
        tracer.enabled = false;
        tracer
    }

    /// 生成唯一的trace_id
    async fn generate_trace_id(&self) -> String {
        let mut counter = self.trace_counter.write().await;
        *counter += 1;
        format!("trace-{:08x}", *counter)
    }

    /// 生成唯一的span_id
    async fn generate_span_id(&self) -> String {
        let mut counter = self.span_counter.write().await;
        *counter += 1;
        format!("span-{:08x}", *counter)
    }

    /// 获取活跃span
    pub async fn get_active_span(&self, span_id: &str) -> Option<TraceSpan> {
        self.active_spans.read().await.get(span_id).cloned()
    }

    /// 获取已完成的spans
    pub async fn get_finished_spans(&self) -> Vec<TraceSpan> {
        self.finished_spans.read().await.clone()
    }

    /// 获取指定trace_id的所有spans
    pub async fn get_spans_by_trace_id(&self, trace_id: &str) -> Vec<TraceSpan> {
        let active = self.active_spans.read().await;
        let finished = self.finished_spans.read().await;
        
        let mut spans = Vec::new();
        
        // 添加活跃spans
        for span in active.values() {
            if span.trace_id == trace_id {
                spans.push(span.clone());
            }
        }
        
        // 添加已完成spans
        for span in finished.iter() {
            if span.trace_id == trace_id {
                spans.push(span.clone());
            }
        }
        
        spans
    }

    /// 清除所有数据
    pub async fn clear_all(&self) {
        self.active_spans.write().await.clear();
        self.finished_spans.write().await.clear();
        *self.statistics.write().await = TraceStatistics {
            total_spans: 0,
            active_spans: 0,
            finished_spans: 0,
            avg_span_duration_ms: 0.0,
            sampling_rate: 1.0,
        };
    }

    /// 计算span持续时间（毫秒）
    fn calculate_duration(&self, span: &TraceSpan) -> f64 {
        if let Some(finish_time) = span.finish_time {
            if let (Ok(start), Ok(finish)) = (
                span.start_time.duration_since(UNIX_EPOCH),
                finish_time.duration_since(UNIX_EPOCH)
            ) {
                return (finish.as_millis() - start.as_millis()) as f64;
            }
        }
        0.0
    }
}

impl Default for MockDistributedTracer {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl DistributedTracer for MockDistributedTracer {
    async fn start_span(&self, operation_name: &str, parent_span: Option<&TraceSpan>) -> TracingResult<TraceSpan> {
        if !self.enabled {
            return Err(Box::new(TracingError::InvalidOperation("追踪器已禁用".to_string())));
        }

        if operation_name.is_empty() {
            return Err(Box::new(TracingError::InvalidData("操作名称不能为空".to_string())));
        }

        let trace_id = match parent_span {
            Some(parent) => parent.trace_id.clone(),
            None => self.generate_trace_id().await,
        };

        let span_id = self.generate_span_id().await;
        let parent_span_id = parent_span.map(|p| p.span_id.clone());

        let span = TraceSpan {
            trace_id,
            span_id: span_id.clone(),
            parent_span_id,
            operation_name: operation_name.to_string(),
            start_time: SystemTime::now(),
            finish_time: None,
            tags: HashMap::new(),
            logs: Vec::new(),
        };

        // 添加到活跃spans
        self.active_spans.write().await.insert(span_id, span.clone());

        // 更新统计信息
        let mut stats = self.statistics.write().await;
        stats.total_spans += 1;
        stats.active_spans = self.active_spans.read().await.len() as u64;

        Ok(span)
    }

    async fn finish_span(&self, span: &mut TraceSpan) -> TracingResult<()> {
        if !self.enabled {
            return Err(Box::new(TracingError::InvalidOperation("追踪器已禁用".to_string())));
        }

        if span.finish_time.is_some() {
            return Err(Box::new(TracingError::InvalidOperation("Span已经完成".to_string())));
        }

        // 设置完成时间
        span.finish_time = Some(SystemTime::now());

        // 从活跃spans中移除
        let mut active_spans = self.active_spans.write().await;
        active_spans.remove(&span.span_id);

        // 添加到已完成spans
        self.finished_spans.write().await.push(span.clone());

        // 更新统计信息
        let mut stats = self.statistics.write().await;
        stats.active_spans = active_spans.len() as u64;
        stats.finished_spans += 1;

        // 更新平均持续时间
        let duration = self.calculate_duration(span);
        if stats.finished_spans == 1 {
            stats.avg_span_duration_ms = duration;
        } else {
            stats.avg_span_duration_ms = 
                (stats.avg_span_duration_ms * (stats.finished_spans - 1) as f64 + duration) 
                / stats.finished_spans as f64;
        }

        Ok(())
    }

    async fn add_span_tag(&self, span: &mut TraceSpan, key: &str, value: &str) -> TracingResult<()> {
        if !self.enabled {
            return Err(Box::new(TracingError::InvalidOperation("追踪器已禁用".to_string())));
        }

        if key.is_empty() {
            return Err(Box::new(TracingError::InvalidData("标签键不能为空".to_string())));
        }

        span.tags.insert(key.to_string(), value.to_string());

        // 同时更新活跃spans中的数据
        if let Some(active_span) = self.active_spans.write().await.get_mut(&span.span_id) {
            active_span.tags.insert(key.to_string(), value.to_string());
        }

        Ok(())
    }

    async fn add_span_log(&self, span: &mut TraceSpan, message: &str, fields: &HashMap<String, String>) -> TracingResult<()> {
        if !self.enabled {
            return Err(Box::new(TracingError::InvalidOperation("追踪器已禁用".to_string())));
        }

        let log = SpanLog {
            timestamp: SystemTime::now(),
            message: message.to_string(),
            fields: fields.clone(),
        };

        span.logs.push(log.clone());

        // 同时更新活跃spans中的数据
        if let Some(active_span) = self.active_spans.write().await.get_mut(&span.span_id) {
            active_span.logs.push(log);
        }

        Ok(())
    }

    async fn export_traces(&self, format: TraceExportFormat) -> TracingResult<String> {
        if !self.enabled {
            return Err(Box::new(TracingError::InvalidOperation("追踪器已禁用".to_string())));
        }

        let finished_spans = self.finished_spans.read().await;

        match format {
            TraceExportFormat::Json => {
                serde_json::to_string_pretty(&*finished_spans)
                    .map_err(|e| Box::new(TracingError::ExportError(e.to_string())) as Box<dyn std::error::Error + Send + Sync>)
            }
            TraceExportFormat::Jaeger => {
                // 简化的Jaeger格式
                let jaeger_data = serde_json::json!({
                    "data": [{
                        "traceID": finished_spans.first().map(|s| &s.trace_id).unwrap_or(&String::new()),
                        "spans": finished_spans.iter().map(|span| {
                            serde_json::json!({
                                "traceID": span.trace_id,
                                "spanID": span.span_id,
                                "parentSpanID": span.parent_span_id,
                                "operationName": span.operation_name,
                                "startTime": span.start_time.duration_since(UNIX_EPOCH).unwrap_or_default().as_micros(),
                                "duration": span.finish_time
                                    .and_then(|ft| ft.duration_since(span.start_time).ok())
                                    .map(|d| d.as_micros())
                                    .unwrap_or(0),
                                "tags": span.tags.iter().map(|(k, v)| serde_json::json!({
                                    "key": k,
                                    "value": v
                                })).collect::<Vec<_>>(),
                                "logs": span.logs.iter().map(|log| serde_json::json!({
                                    "timestamp": log.timestamp.duration_since(UNIX_EPOCH).unwrap_or_default().as_micros(),
                                    "fields": log.fields.iter().map(|(k, v)| serde_json::json!({
                                        "key": k,
                                        "value": v
                                    })).collect::<Vec<_>>()
                                })).collect::<Vec<_>>()
                            })
                        }).collect::<Vec<_>>()
                    }]
                });
                
                serde_json::to_string_pretty(&jaeger_data)
                    .map_err(|e| Box::new(TracingError::ExportError(e.to_string())) as Box<dyn std::error::Error + Send + Sync>)
            }
            TraceExportFormat::OpenTelemetry => {
                // 简化的OpenTelemetry格式
                let otel_data = serde_json::json!({
                    "resourceSpans": [{
                        "scopeSpans": [{
                            "spans": finished_spans.iter().map(|span| serde_json::json!({
                                "traceId": span.trace_id,
                                "spanId": span.span_id,
                                "parentSpanId": span.parent_span_id,
                                "name": span.operation_name,
                                "startTimeUnixNano": span.start_time.duration_since(UNIX_EPOCH).unwrap_or_default().as_nanos(),
                                "endTimeUnixNano": span.finish_time
                                    .and_then(|ft| ft.duration_since(UNIX_EPOCH).ok())
                                    .map(|d| d.as_nanos())
                                    .unwrap_or(0),
                                "attributes": span.tags.iter().map(|(k, v)| serde_json::json!({
                                    "key": k,
                                    "value": {"stringValue": v}
                                })).collect::<Vec<_>>()
                            })).collect::<Vec<_>>()
                        }]
                    }]
                });

                serde_json::to_string_pretty(&otel_data)
                    .map_err(|e| Box::new(TracingError::ExportError(e.to_string())) as Box<dyn std::error::Error + Send + Sync>)
            }
        }
    }

    async fn get_trace_statistics(&self) -> TracingResult<TraceStatistics> {
        if !self.enabled {
            return Err(Box::new(TracingError::InvalidOperation("追踪器已禁用".to_string())));
        }

        let stats = self.statistics.read().await.clone();
        Ok(stats)
    }
}

/// 基于UUID的真实追踪器
/// 
/// 生产环境使用的追踪器实现，使用UUID生成唯一ID
pub struct UuidDistributedTracer {
    /// 活跃的spans
    active_spans: Arc<RwLock<HashMap<String, TraceSpan>>>,
    /// 已完成的spans
    finished_spans: Arc<RwLock<Vec<TraceSpan>>>,
    /// 统计信息
    statistics: Arc<RwLock<TraceStatistics>>,
    /// 是否启用
    enabled: bool,
    /// 采样率
    sampling_rate: f64,
}

impl UuidDistributedTracer {
    /// 创建新的UUID追踪器
    pub fn new() -> Self {
        Self {
            active_spans: Arc::new(RwLock::new(HashMap::new())),
            finished_spans: Arc::new(RwLock::new(Vec::new())),
            statistics: Arc::new(RwLock::new(TraceStatistics {
                total_spans: 0,
                active_spans: 0,
                finished_spans: 0,
                avg_span_duration_ms: 0.0,
                sampling_rate: 1.0,
            })),
            enabled: true,
            sampling_rate: 1.0,
        }
    }

    /// 创建带采样率的追踪器
    pub fn with_sampling_rate(sampling_rate: f64) -> Self {
        let mut tracer = Self::new();
        tracer.sampling_rate = sampling_rate.clamp(0.0, 1.0);
        tracer.statistics.clone().try_write().unwrap().sampling_rate = tracer.sampling_rate;
        tracer
    }

    /// 生成唯一的trace_id
    fn generate_trace_id(&self) -> String {
        Uuid::new_v4().to_string().replace('-', "")
    }

    /// 生成唯一的span_id
    fn generate_span_id(&self) -> String {
        Uuid::new_v4().to_string().replace('-', "")[0..16].to_string()
    }

    /// 判断是否应该采样
    fn should_sample(&self) -> bool {
        if self.sampling_rate >= 1.0 {
            return true;
        }
        if self.sampling_rate <= 0.0 {
            return false;
        }
        
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};
        
        let mut hasher = DefaultHasher::new();
        SystemTime::now().hash(&mut hasher);
        let hash = hasher.finish();
        let random = (hash as f64) / (u64::MAX as f64);
        random < self.sampling_rate
    }
}

impl Default for UuidDistributedTracer {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl DistributedTracer for UuidDistributedTracer {
    async fn start_span(&self, operation_name: &str, parent_span: Option<&TraceSpan>) -> TracingResult<TraceSpan> {
        if !self.enabled {
            return Err(Box::new(TracingError::InvalidOperation("追踪器已禁用".to_string())));
        }

        if operation_name.is_empty() {
            return Err(Box::new(TracingError::InvalidData("操作名称不能为空".to_string())));
        }

        // 检查采样率
        if !self.should_sample() {
            return Err(Box::new(TracingError::InvalidOperation("被采样器跳过".to_string())));
        }

        let trace_id = match parent_span {
            Some(parent) => parent.trace_id.clone(),
            None => self.generate_trace_id(),
        };

        let span_id = self.generate_span_id();
        let parent_span_id = parent_span.map(|p| p.span_id.clone());

        let span = TraceSpan {
            trace_id,
            span_id: span_id.clone(),
            parent_span_id,
            operation_name: operation_name.to_string(),
            start_time: SystemTime::now(),
            finish_time: None,
            tags: HashMap::new(),
            logs: Vec::new(),
        };

        // 添加到活跃spans
        self.active_spans.write().await.insert(span_id, span.clone());

        // 更新统计信息
        let mut stats = self.statistics.write().await;
        stats.total_spans += 1;
        stats.active_spans = self.active_spans.read().await.len() as u64;

        Ok(span)
    }

    async fn finish_span(&self, span: &mut TraceSpan) -> TracingResult<()> {
        if !self.enabled {
            return Err(Box::new(TracingError::InvalidOperation("追踪器已禁用".to_string())));
        }

        if span.finish_time.is_some() {
            return Err(Box::new(TracingError::InvalidOperation("Span已经完成".to_string())));
        }

        // 设置完成时间
        span.finish_time = Some(SystemTime::now());

        // 从活跃spans中移除
        let mut active_spans = self.active_spans.write().await;
        active_spans.remove(&span.span_id);

        // 添加到已完成spans
        self.finished_spans.write().await.push(span.clone());

        // 更新统计信息
        let mut stats = self.statistics.write().await;
        stats.active_spans = active_spans.len() as u64;
        stats.finished_spans += 1;

        Ok(())
    }

    async fn add_span_tag(&self, span: &mut TraceSpan, key: &str, value: &str) -> TracingResult<()> {
        if !self.enabled {
            return Err(Box::new(TracingError::InvalidOperation("追踪器已禁用".to_string())));
        }

        if key.is_empty() {
            return Err(Box::new(TracingError::InvalidData("标签键不能为空".to_string())));
        }

        span.tags.insert(key.to_string(), value.to_string());

        // 同时更新活跃spans中的数据
        if let Some(active_span) = self.active_spans.write().await.get_mut(&span.span_id) {
            active_span.tags.insert(key.to_string(), value.to_string());
        }

        Ok(())
    }

    async fn add_span_log(&self, span: &mut TraceSpan, message: &str, fields: &HashMap<String, String>) -> TracingResult<()> {
        if !self.enabled {
            return Err(Box::new(TracingError::InvalidOperation("追踪器已禁用".to_string())));
        }

        let log = SpanLog {
            timestamp: SystemTime::now(),
            message: message.to_string(),
            fields: fields.clone(),
        };

        span.logs.push(log.clone());

        // 同时更新活跃spans中的数据
        if let Some(active_span) = self.active_spans.write().await.get_mut(&span.span_id) {
            active_span.logs.push(log);
        }

        Ok(())
    }

    async fn export_traces(&self, format: TraceExportFormat) -> TracingResult<String> {
        if !self.enabled {
            return Err(Box::new(TracingError::InvalidOperation("追踪器已禁用".to_string())));
        }

        let finished_spans = self.finished_spans.read().await;

        match format {
            TraceExportFormat::Json => {
                serde_json::to_string_pretty(&*finished_spans)
                    .map_err(|e| Box::new(TracingError::ExportError(e.to_string())) as Box<dyn std::error::Error + Send + Sync>)
            }
            _ => {
                // 其他格式暂时返回JSON
                serde_json::to_string_pretty(&*finished_spans)
                    .map_err(|e| Box::new(TracingError::ExportError(e.to_string())) as Box<dyn std::error::Error + Send + Sync>)
            }
        }
    }

    async fn get_trace_statistics(&self) -> TracingResult<TraceStatistics> {
        if !self.enabled {
            return Err(Box::new(TracingError::InvalidOperation("追踪器已禁用".to_string())));
        }

        let stats = self.statistics.read().await.clone();
        Ok(stats)
    }
}